load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/gsn_code.ncl"
load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/gsn_csm.ncl"
load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/contributed.ncl"
load "/WORK/jiangj/Function/Transfer_dimension_name.ncl"

begin 

;=====================================input:u,v,w,q,ps,E,P for each month=================
do m=5,8

print(tostring(m+1))

fread = addfile("data/reana/JRA55/JRA55.ugrd.mon.195801-202212.grb","r")
u     = fread->UGRD_GDS0_ISBL_S123(m::12,{1000:100},{10:50},{60:110});m/s
u     = Transter_dimension4_name(u)

fread = addfile("data/reana/JRA55/JRA55.vgrd.mon.195801-202212.grb","r")
v     = fread->VGRD_GDS0_ISBL_S123(m::12,{1000:100},{10:50},{60:110});m/s
v     = Transter_dimension4_name(v)

fread = addfile("data/reana/JRA55/JRA55.vvel.mon.195801-202212.grb","r")
w     = fread->VVEL_GDS0_ISBL_S123(m::12,{1000:100},{10:50},{60:110});Pa/s
w     = Transter_dimension4_name(w)

fread = addfile("data/reana/JRA55/JRA55.spfh.mon.195801-202212.grb","r")
q     = fread->SPFH_GDS0_ISBL_S123(m::12,{1000:100},{10:50},{60:110});kg/kg
q     = Transter_dimension4_name(q)

fread = addfile("data/reana/JRA55/JRA55.pres.mon.195801-202212.grb","r")
ps    = fread->PRES_GDS0_SFC_S123(m::12,{10:50},{60:110});Pa
ps    = ps/100.;hPa 
ps    = Transter_dimension3_name(ps)

fread = addfile("data/reana/JRA55/JRA55.evp.mon.195801-202212.grb","r")
E     = fread->EVP_GDS0_SFC_S130(m::12,{10:50},{60:110});mm/day
E     = Transter_dimension3_name(E)

fread = addfile("data/reana/JRA55/JRA55.tprat.mon.195801-202212.grb","r")
pr    = fread->TPRAT_GDS0_SFC_S130(m::12,{10:50},{60:110});mm/day
pr    = Transter_dimension3_name(pr)

ub = dim_avg_n_Wrap(u,0)
vb = dim_avg_n_Wrap(v,0)
wb = dim_avg_n_Wrap(w,0)
qb = dim_avg_n_Wrap(q,0)

plev  = u&lev
lat   = u&lat
lon   = u&lon

nlat = dimsizes(lat)
nlon = dimsizes(lon)
dlat = abs(lat(2) - lat(1)) * 0.0174533
dlon = abs(lon(2) - lon(1)) * 0.0174533

;****************************************************************
;wind: m/s; specific humidy Kg/Kg; ps hPa; omega Pa/s; lev hPa**
;*******************Water vapor transport flux volumn***********
uq = u
vq = v
uq = u*q
vq = v*q
linlog = 1
g = 9.8
pbot = max(plev)
ptop = min(plev)
UQ_vint = 1./g * 100 * vibeta (plev,uq(time|:,lat|:,lon|:,lev|:),linlog,ps,pbot,ptop) ; kg/(m*s) 
VQ_vint = 1./g * 100 * vibeta (plev,vq(time|:,lat|:,lon|:,lev|:),linlog,ps,pbot,ptop)
copy_VarCoords(ps,UQ_vint)
copy_VarCoords(ps,VQ_vint)
UQ_vint@units = "kg m-1 s-1"
VQ_vint@units = "kg m-1 s-1"

;*******************Water vapor budget volumn****************
; P = -<delt(Vq)> + E + residual  ***  -<delt(Vq)> = -wdqdp-udqdx-vdqdy
; -wdqdp
wdqdp = w;
dqdp  = center_finite_diff_n(q, plev, False, 0, 1)
wdqdp = w * dqdp

; -udqdx 
udqdx = u ;(time|:,lat|:,lon|:,lev|:)
do  i = 0, nlat-1
    dx = 6378388. *cos(0.0174533 *lat(i)) * dlon
    udqdx(:,:,i,:) = center_finite_diff_n(q(:,:,i,:), dx, False, 0, 2)
end do
udqdx = u * udqdx

; -vdqdy
vdqdy = v ;(time|:,lat|:,lon|:,lev|:)
do i = 0,nlon-1
   dy = 6378388.*dlat
   vdqdy(:,:,:,i) = center_finite_diff_n(q(:,:,:,i), dy, False, 0, 2)
end do
vdqdy = v * vdqdy

udqdx_vint = -1./g * 100 * 86400 * vibeta (plev,udqdx(time|:,lat|:,lon|:,lev|:),linlog,ps,pbot,ptop) ;to mm/day
vdqdy_vint = -1./g * 100 * 86400 * vibeta (plev,vdqdy(time|:,lat|:,lon|:,lev|:),linlog,ps,pbot,ptop) ;to mm/day
wdqdp_vint = -1./g * 86400 * vibeta (plev,wdqdp(time|:,lat|:,lon|:,lev|:),linlog,ps,pbot,ptop) ;to mm/day

copy_VarCoords(ps,udqdx_vint)
copy_VarCoords(ps,vdqdy_vint)
copy_VarCoords(ps,wdqdp_vint)

; -wdqdp ~= -wbdqpdp -wpdqbdp (qp = q - q_mean; wp = w -w_mean)
wbdqpdp = q
wpdqbdp = w 
qp = q - conform(q, qb, (/1,2,3/))
wp = w - conform(w, wb, (/1,2,3/))
dqpdp   = center_finite_diff_n(qp ,plev,False,0, 1)
dqbdp   = center_finite_diff_n(qb ,plev,False,0, 0)
wbdqpdp = conform(w, wb, (/1,2,3/)) * dqpdp
wpdqbdp = wp * conform(w, dqbdp, (/1,2,3/))

Ther_w = -1./g * 86400 * vibeta (plev,wbdqpdp(time|:,lat|:,lon|:,lev|:),linlog,ps,pbot,ptop) ;to mm/day
Dyna_w = -1./g * 86400 * vibeta (plev,wpdqbdp(time|:,lat|:,lon|:,lev|:),linlog,ps,pbot,ptop) ;to mm/day

copy_VarCoords(ps,Ther_w)
copy_VarCoords(ps,Dyna_w)

;-udqdx ~ = -ubdqpdx-updqbdx 
up = u - conform(u, ub, (/1,2,3/))

ubdqpdx = q ;(time|:,lat|:,lon|:,lev|:)
do  i = 0, nlat-1
    dx = 6378388. *cos(0.0174533 *lat(i)) * dlon
    ubdqpdx(:,:,i,:) = center_finite_diff_n(qp(:,:,i,:), dx, False, 0, 2)
end do
ubdqpdx = conform(u,ub,(/1,2,3/))*ubdqpdx 

updqbdx = u 
do  i = 0, nlat-1
    dx = 6378388. *cos(0.0174533 *lat(i)) * dlon
    updqbdx(:,:,i,:) = conform(u(:,:,0,:), center_finite_diff_n(qb(:,i,:), dx, False, 0, 1), (/1,2/))
end do
updqbdx = up*updqbdx 

Ther_u = -1./g * 100 * 86400 * vibeta (plev,ubdqpdx(time|:,lat|:,lon|:,lev|:),linlog,ps,pbot,ptop)
Dyna_u = -1./g * 100 * 86400 * vibeta (plev,updqbdx(time|:,lat|:,lon|:,lev|:),linlog,ps,pbot,ptop)

copy_VarCoords(ps,Ther_u)
copy_VarCoords(ps,Dyna_u)

;-udqdx ~ = -ubdqpdx-updqbdx 
vp = v - conform(v, vb, (/1,2,3/))

vbdqpdy = q ;(time|:,lat|:,lon|:,lev|:)
do i = 0,nlon-1
   dy = 6378388.*dlat
   vbdqpdy(:,:,:,i) = center_finite_diff_n(qp(:,:,:,i), dy, False, 0, 2)
end do
vbdqpdy = conform(v,vb,(/1,2,3/)) * vbdqpdy

vpdqbdy = v ;(time|:,lat|:,lon|:,lev|:)
do i = 0,nlon-1
   dy = 6378388.*dlat
   vpdqbdy(:,:,:,i) = conform(v(:,:,:,0),center_finite_diff_n(qb(:,:,i), dy, False, 0, 1),(/1,2/))
end do
vpdqbdy = vp * vpdqbdy

Ther_v = -1./g * 100 * 86400 * vibeta (plev,vbdqpdy(time|:,lat|:,lon|:,lev|:),linlog,ps,pbot,ptop) ;to mm/day
Dyna_v = -1./g * 100 * 86400 * vibeta (plev,vpdqbdy(time|:,lat|:,lon|:,lev|:),linlog,ps,pbot,ptop) ;to mm/day

copy_VarCoords(ps,Ther_v)
copy_VarCoords(ps,Dyna_v)

;=================================
updqpdx = q ;(time|:,lat|:,lon|:,lev|:)
do  i = 0, nlat-1
    dx = 6378388. *cos(0.0174533 *lat(i)) * dlon
    updqpdx(:,:,i,:) = center_finite_diff_n(qp(:,:,i,:), dx, False, 0, 2)
end do
updqpdx = up*updqpdx 

vpdqpdy = q ;(time|:,lat|:,lon|:,lev|:)
do i = 0,nlon-1
   dy = 6378388.*dlat
   vpdqpdy(:,:,:,i) = center_finite_diff_n(qp(:,:,:,i), dy, False, 0, 2)
end do
vpdqpdy = vp * vpdqpdy

wpdqpdp = q
wpdqpdp = wp * dqpdp

Res_u = -1./g * 100 * 86400 * vibeta (plev,updqpdx(time|:,lat|:,lon|:,lev|:),linlog,ps,pbot,ptop) ;to mm/day
Res_v = -1./g * 100 * 86400 * vibeta (plev,vpdqpdy(time|:,lat|:,lon|:,lev|:),linlog,ps,pbot,ptop) ;to mm/day
Res_w = -1./g * 86400 * vibeta (plev,wpdqpdp(time|:,lat|:,lon|:,lev|:),linlog,ps,pbot,ptop) ;to mm/day

copy_VarCoords(ps,Res_w)
copy_VarCoords(ps,Res_u)
copy_VarCoords(ps,Res_v)

;==================================output=========================================
udqdx_vint@units = "mm/day"
vdqdy_vint@units = "mm/day"
wdqdp_vint@units = "mm/day"
Ther_u@units     = "mm/day"
Dyna_u@units     = "mm/day"
Ther_v@units     = "mm/day"
Dyna_v@units     = "mm/day"
Ther_w@units     = "mm/day"
Dyna_w@units     = "mm/day"
Res_u@units      = "mm/day"
Res_v@units      = "mm/day"
Res_w@units      = "mm/day"
E@units          = "mm/day"
pr@units         = "mm/day"

out_fname = "data/MB_1958-2022_0"+tostring(m+1)+"_JRA55.nc"
system("rm -rf "+out_fname)
fout = addfile(out_fname,"c")
fout->UQ_vint  = UQ_vint
fout->VQ_vint  = VQ_vint
fout->U_budget = udqdx_vint
fout->V_budget = vdqdy_vint
fout->W_budget = wdqdp_vint
fout->W_Ther   = Ther_w
fout->W_Dyna   = Dyna_w
fout->U_Ther   = Ther_u
fout->U_Dyna   = Dyna_u
fout->V_Ther   = Ther_v
fout->V_Dyna   = Dyna_v
fout->U_Res    = Res_u
fout->V_Res    = Res_v
fout->W_Res    = Res_w
fout->E        = E
fout->P        = pr

end do 

end

